from keras.models import Model
from keras.layers import Embedding, Bidirectional, LSTM, Dense, Dropout, Input
from keras_contrib.layers.crf import CRF, crf_loss, crf_viterbi_accuracy

from DataProcess.process_data import DataProcess
from sklearn.metrics import f1_score
import numpy as np
from keras.utils.vis_utils import plot_model
from keras.models import load_model
from keras import backend as K




class BILSTMCRF():
    def __init__(self,
                 vocab_size: int,
                 n_class: int,
                 embedding_dim: int = 128,
                 rnn_units: int = 128,
                 drop_rate: float = 0.5,
                 ):
        self.vocab_size = vocab_size
        self.n_class = n_class
        self.embedding_dim = embedding_dim
        self.rnn_units = rnn_units
        self.drop_rate = drop_rate

    def creat_model(self):
        inputs = Input(shape=(None,))
        x = Embedding(input_dim=self.vocab_size, output_dim=self.embedding_dim)(inputs)
        x = Bidirectional(LSTM(units=self.rnn_units, return_sequences=True))(x)
        x = Dropout(self.drop_rate)(x)
        x = Dense(self.n_class)(x)
        self.crf = CRF(self.n_class, sparse_target=False)
        x = self.crf(x)
        self.model = Model(inputs=inputs, outputs=x)
        self.model.summary()
        self.compile()
        return self.model

    def compile(self):
        self.model.compile('adam',
                           loss=self.crf.loss_function,
                           metrics=[self.crf.accuracy])


# 用于BiLSTM_CRF的模型训练
def trainBiLSTMCRF(datatype):
    dp = DataProcess(data_type=datatype)
    train_data, train_label, test_data, test_label = dp.get_data(one_hot=True)
    print(test_data)
    # lstm_crf = BILSTMCRF(vocab_size=dp.vocab_size, n_class=7)
    # lstm_crf.creat_model()
    # model = lstm_crf.model
    # plot_model(model, to_file='picture/BILSTM_CRF.png', show_shapes=True)
    # # exit()
    #
    # model.fit(train_data, train_label, batch_size=64, epochs=100,
    #           validation_data=[test_data, test_label])
    # model.save("model/BiLSTM_CRF.h5")
    # print(datatype + "_BiLSTMCRF模型保存成功")
    # exit()


# 获得保存的模型，测试集进行测试
def modelValid(datatype):
    dp = DataProcess(data_type=datatype)
    train_data, train_label, test_data, test_label = dp.get_data(one_hot=True)
    model = load_model("model/BiLSTM_CRF.h5",
                       custom_objects={"CRF": CRF, 'crf_loss': crf_loss, 'crf_viterbi_accuracy': crf_viterbi_accuracy})

    # 评估模型
    score = model.evaluate(x=test_data, y=test_label, verbose=0)
    print("Loss on test set：", score[0])  # 0.08888126758734385   0.030318434461951257
    print("Accuracy on test set:", score[1])  # 0.9666399987538655   0.989760007540385

    # 测试结果
    y = model.predict(test_data)

    # 对比测试数据的tag
    label_indexs = []
    predict_indexs = []
    num2tag = dp.num2tag()
    i2w = dp.i2w()
    texts = []
    texts.append(f"字符\t预测tag\t原tag\n")
    for i, x_line in enumerate(test_data):
        for j, index in enumerate(x_line):
            if index != 0:
                char = i2w.get(index, ' ')
                t_line = y[i]
                t_index = np.argmax(t_line[j])
                tag = num2tag.get(t_index, 'O')
                predict_indexs.append(t_index)

                t_line = test_label[i]
                t_index = np.argmax(t_line[j])
                org_tag = num2tag.get(t_index, 'O')
                label_indexs.append(t_index)

                texts.append(f"{char}\t{tag}\t{org_tag}\n")
        texts.append('\n')

    f1score = f1_score(label_indexs, predict_indexs, average='macro')
    print(f"f1score:{f1score}")

    # exit()
    # 预测结果
    with open('./predict/bilstmcrf_pre_data.txt', 'w', encoding='utf-8') as f:
        f.write("".join(texts))


# 传入要预测的字符串以及字符串所属数据集的类型，结果得到预测出来的标签列表
def predict(datatype, rawpredictdata, select):
    dp = DataProcess(data_type=datatype)
    if select == 1:
        model = load_model(r"F:\develop_code\py_code\demo\demo\model\BiLSTM_CRF.h5",
                           custom_objects={"CRF": CRF, 'crf_loss': crf_loss,
                                           'crf_viterbi_accuracy': crf_viterbi_accuracy})
    else:
        model = load_model("model/BiLSTM_CRF_epoch50.h5",
                           custom_objects={"CRF": CRF, 'crf_loss': crf_loss,
                                           'crf_viterbi_accuracy': crf_viterbi_accuracy})
    # 对原始字符串数据进行处理形成BiLSTMCRF模型的输入
    predictdata = dp.processraw(rawpredictdata)
    # print("原始输入数据：")
    # print(predictdata)
    # print("原始输入数据的格式")
    # print(predictdata.shape)  (1, 100)
    y = model.predict(predictdata)
    K.clear_session()
    # print("测试结果")
    # print(y)
    # print(y.shape) (1, 100, 7)

    # 数字对应的标签 如 1：B-PER
    num2tag = dp.num2tag()
    # 索引对应的单词
    i2w = dp.i2w()
    texts = []
    charres = []
    tagres = []
    texts.append(f"字符\t预测tag\n")
    # enumerate()返回的是索引（从0开始）+列表中的元素
    # x_line表示的是一个句子，长度是100[], predictdata的格式是[[],[],[],...],每个内部[]的长度都是100
    for i, x_line in enumerate(predictdata):
        # 接下来的这个循环是访问长度为100的句子当中的每个单词
        for j, index in enumerate(x_line):
            if index != 0:
                # index 对应的单词 char
                char = i2w.get(index, ' ')
                # 每个单词对应的标签的表示都是一个长度为7的列表 ，t_line表示的是一个句子的标签表示
                t_line = y[i]  # y.shape= （1， 100， 7） ，因为只有一个句子，所以第一个维度是1
                t_index = np.argmax(t_line[j])  # np.argmax() 返回的是列表中的最大值的索引
                # id 对应的标签 tag
                tag = num2tag.get(t_index, 'O')
                charres.append(char)
                tagres.append(tag)
                texts.append(f"{char}\t{tag}\n")
        texts.append('\n')
    # for i in texts:
    #     print(i)
    return charres, tagres


# 将最终预测出来的两列结果（句子，句子对应的标签，长度都是100）中的标签结果提取出来, select使用不同的ephco训练出来的模型进行预测
def resultprocess(datatype, rawpredictdata, select):
    # charres, tagres是两个列表
    charres, tagres = predict(datatype, rawpredictdata, select)
    phe = ""
    cau = ""
    mnt = ""
    flag = False
    # 如何取出以B-开头的
    for s, t in zip(charres, tagres):
        if t == 'O':
            flag = False
        if t in ('B-PHE', 'I-PHE'):
            if t == 'B-PHE':
                flag = True
                phe += s
            elif t == 'I-PHE':
                if flag:
                    phe += s
                else:
                    # 此时表示找到的是I-开始的并不是以B-开始的
                    continue
        if t in ('B-CAU', 'I-CAU'):
            if t == 'B-CAU':
                flag = True
                cau += s
            elif t == 'I-CAU':
                if flag:
                    cau += s
                else:
                    continue
        if t in ('B-MNT', 'I-MNT'):
            if t == 'B-MNT':
                flag = True
                mnt += s
            elif t == 'I-MNT':
                if flag:
                    mnt += s
                else:
                    continue
        # print(s + "," + t)
        # for s, t in zip(charres, tagres):
        #     if t in ('B-PHE', 'I-PHE'):
        #         phe += ' ' + s if (t == 'B-PHE') else s
        #     if t in ('B-CAU', 'I-CAU'):
        #         cau += ' ' + s if (t == 'B-CAU') else s
        #     if t in ('B-MNT', 'I-MNT'):
        #         mnt += ' ' + s if (t == 'B-MNT') else s
        #     print(s + "," + t)
    if phe != "":
        return dict(phe=phe)
    elif cau != "":
        return dict(cau=cau)
    elif mnt !="":
        return dict(mnt=mnt)
    return ""



if __name__ == '__main__':
    # 模型训练,先后采用epchos = 50, 100
    trainBiLSTMCRF('rbfd')
    # 验证模型准确率
    # modelValid('rbfd')
    # datatype = 'rbfd'

    # 识别故障现象PHE    成功
    # rawpredictdata = "出现塑性变形的原因是什么"
    # rawpredictdata = "出现点蚀的原因是什么"
    # rawpredictdata = "微动磨损的原因是什么"

    # 识别故障原因原因CAU  成功
    # rawpredictdata = "润滑剂选用不当会引发什么故障"
    # rawpredictdata = "轴承内部含有酸会引发什么故障"

    # 识别故障维护措施MNT   部分成功
    # rawpredictdata = "提高耐腐蚀能力可以防止哪些现象的发生"
    # rawpredictdata = "表面改性可以防止哪些现象的发生" # 两次都失败
    # rawpredictdata = "提高过盈量可以防止哪些现象的发生"

    # res = resultprocess(datatype, rawpredictdata, 1)
    # for key, value in res.items():
    #     print(key, "->", value)

    # 使用不同epcho得到的模型进行预测
    # print("第二次预测")
    # resutprocess(datatype, rawpredictdata, 2)
